Skip to content

Ask SLURM for what the job holds when the fit is launched from the login node (#642) - #645

Merged
wshlavacek merged 1 commit into
mainfrom
fix/642-srun-login-node-cpu-count
Aug 22, 2026
Merged

Ask SLURM for what the job holds when the fit is launched from the login node (#642)#645
wshlavacek merged 1 commit into
mainfrom
fix/642-srun-login-node-cpu-count

Conversation

@wshlavacek

Copy link
Copy Markdown
Collaborator

Fixes #642.

What was wrong

On a cluster where salloc returns a shell on the login node while the allocation is held on a compute node — which is the shell PyBNF is meant to run in, since it is the one that holds the allocation — the srun launcher stopped before a single worker started:

srun: error: Unable to create step for job NNNNN: More processors requested than permitted

Cluster.cpus_per_node prefers $SLURM_CPUS_ON_NODE, which SLURM sets only inside a job step running on an allocated node. On the login node it is simply absent, so the count fell through to dask.system.CPU_COUNT and then to the machine's processor count — both of which describe the login node, which is not in the allocation and is usually several times larger than what the job holds. A job granted 20 CPUs was sized as though it held 128, and --cpus-per-task 128 is a request SLURM refuses outright.

The fix

Two parts, both in pybnf/cluster.py. ADR-0125 records the reasoning.

1. The job's own per-node list is a source of the single count. cpus_per_node now consults $SLURM_JOB_CPUS_PER_NODE — the per-node list SLURM publishes for the job, which is set correctly in a login-node shell — between $SLURM_CPUS_ON_NODE and the two machine-level numbers. Its smallest entry is taken, because this single number both sizes a pool started on every machine and is what one srun step asks for on every machine in it: asking for less than a machine holds costs speed, asking for more than the smallest holds is refused. Inside the allocation nothing changes — $SLURM_CPUS_ON_NODE is still preferred where SLURM sets it. This also fixes the same login-node exposure on the SSH launcher (-t slurm), which reads the same count.

2. The layout hands its counts on rather than having them re-derived (the fix suggested in the issue). srun_worker_layout already reads the per-node list before choosing which command to build; the equal-size path now passes that count — and the phrase naming its source, for the log — into srun_worker_command, which derives one itself only when no caller supplied it. So a $SLURM_CPUS_ON_NODE that does not describe this job, including one exported by hand as the workaround for this bug, no longer sizes a later run.

The srun argument list is otherwise untouched: a launch from inside the allocation builds the command it built before, and the existing regression test that pins that argv still passes.

Verification

The constructed srun argument list is the oracle, as it was for #614 and #617. Eight new tests; all seven that target the defect fail on main and pass here:

  • cpus_per_node — the job list is read when the step variable is empty; a mixed list reduces to its smallest machine; the step variable still wins when both are set; an unreadable or zero list falls through to the local numbers.
  • srun_worker_command — a caller-supplied count is used instead of the environment (and its source reaches the log); it still caps an oversubscribed parallel_count.
  • setup_srun_cluster, end to end in a stood-up login-node environment (no $SLURM_CPUS_ON_NODE, a job list of 20, machine-level numbers of 128) — the default path asks for 20 workers and 20 CPUs rather than 128; an explicit parallel_count = 64 starts all 64 workers but asks for 20 CPUs, so the step is not refused; and the count the layout read is the count the command uses even when $SLURM_CPUS_ON_NODE is stale.
  • per_node_cpus — the two fixes compose: a list that cannot be lined up with the machines still supplies the fallback count rather than the login node's size.

Full local gate green (4564 passed) apart from one unrelated pre-existing failure — test_gradient_sens_fallback.py::test_a_declined_model_is_reported_identically_cold_and_warm, which fails on a clean main too against a bngsim built from upstream main and is filed as #644. Docs updated (docs/cluster.rst, docs/config_keys.rst) and a CHANGELOG entry added.

Also noted

ADR-0124 recorded a limitation on the parallel_count path — on a mixed-size allocation its even split can ask for more CPUs than the smallest machine holds — and asked for its own tracking issue rather than being folded into that change. That is now filed as #643, and this PR deliberately leaves it alone: what changed here is only that the single count it is capped by describes the allocation rather than the machine PyBNF is launched from.

…gin node (#642)

On many clusters salloc returns a shell on the login node while the allocation is
held on a compute node. That shell is where PyBNF is meant to run, because it is the
one that holds the allocation. Started there, the srun launcher stopped before a
single worker came up, with srun reporting that more processors were requested than
permitted.

The count PyBNF sizes a run by came from $SLURM_CPUS_ON_NODE, which SLURM sets only
inside a job step running on an allocated node. On the login node it is absent, so
the count fell through to two numbers that describe the machine doing the asking. On
the login node those describe the login node, which is not in the allocation and is
usually several times larger than what the job holds, and asking SLURM for that many
CPUs per task is a request it refuses.

Cluster.cpus_per_node now consults $SLURM_JOB_CPUS_PER_NODE, the per-node list SLURM
publishes for the job as a whole, between $SLURM_CPUS_ON_NODE and the two
machine-level numbers. That variable is set correctly in a login-node shell. Its
smallest entry is the one taken, because a single number has to be acceptable on
every machine an srun step runs on: asking for fewer CPUs than a machine holds costs
speed, while asking for more than the smallest holds is refused outright. Inside the
allocation nothing changes, since $SLURM_CPUS_ON_NODE is still preferred wherever
SLURM sets it. The SSH launcher reads the same count, so it too stops sizing each
remote machine by the login node.

srun_worker_layout also hands on the per-machine counts it has already read rather
than having srun_worker_command derive the number a second time, which is the fix the
issue suggested. A $SLURM_CPUS_ON_NODE that does not describe this job, including one
exported by hand as the workaround for this bug, no longer sizes a later run. The
argument list is otherwise untouched, so a launch from inside the allocation builds
the command it built before.

The constructed srun command is the oracle, as it was for #614 and #617. Eight new
tests cover the new source and its precedence, a caller-supplied count reaching both
the command and the log, and an end-to-end login-node environment on the default and
parallel_count paths. ADR-0125 records the reasoning; #643 tracks the parallel_count
limitation ADR-0124 deferred.
@wshlavacek
wshlavacek merged commit 338d5d0 into main Aug 22, 2026
9 checks passed
@wshlavacek
wshlavacek deleted the fix/642-srun-login-node-cpu-count branch August 22, 2026 22:01
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

The srun launcher asks for too many CPUs when run from the login node

1 participant